#include "Renderer_PCH.h"
#include "RMaterial.h"

/*--------------------------------------------------------------------------*/

RMaterial::RMaterial(const std::string& name) :
	m_materialName(name),
	m_rasterizerstate(NULL),
	m_depthstencilstate(NULL),
	m_blendstate(NULL),
	m_handle(0),
	m_normalVersion(NULL),
	m_instancedVersion(NULL),
	m_sizeofFconsts(0),
	m_sizeofGconsts(0),
	m_sizeofVconsts(0)
{
	m_normalVersion = this;
}

/*--------------------------------------------------------------------------*/

RMaterial::~RMaterial()
{
	if (m_handle != 0)
	{
		glDeleteProgram(m_handle);
	}
}

/*--------------------------------------------------------------------------*/

const RRasterizerState* RMaterial::GetRasterizerState() const
{
	return m_rasterizerstate;
}

/*--------------------------------------------------------------------------*/

const RDepthstencilState* RMaterial::GetDepthstencilState() const
{
	return m_depthstencilstate;
}

/*--------------------------------------------------------------------------*/

const RBlendState* RMaterial::GetBlendState() const
{
	return m_blendstate;
}

/*--------------------------------------------------------------------------*/

bool RMaterial::CreateProgram( const std::string& vCode, const std::string& gCode, const std::string& fCode)
{
	m_vCode = vCode;
	m_gCode = gCode;
	m_fCode = fCode;

	if (m_handle)
	{
		glDeleteProgram(m_handle);
	}

	m_handle = glCreateProgram();
	__Check(m_handle);

	std::string defines;
	for (auto iter = m_macros.begin(); iter != m_macros.end(); ++iter)
	{
		defines += "#define " + iter->GetName() + " " + iter->GetValue() + "\n";
	}

	uint smp_counter = 0;
	for (auto iter = m_samplers.begin(); iter != m_samplers.end(); ++iter)
	{
		defines += "#define USE_SAMPLER_" + std::to_string(smp_counter) + "\n" + 
				   "#define SAMPLER_" + std::to_string(smp_counter) + "_NAME " + iter->m_name + "\n" +
				   "#define SAMPLER_" + iter->m_name + " " + "SAMPLER_" + std::to_string(smp_counter) + "_NAME\n";

		smp_counter++;
	}

	uint vAttached = AttachShader("VERTEX_SHADER",	GL_VERTEX_SHADER,	defines, vCode);
	uint gAttached = AttachShader("GEOMETRY_SHADER", GL_GEOMETRY_SHADER, defines, gCode);
	uint fAttached = AttachShader("FRAGMENT_SHADER", GL_FRAGMENT_SHADER, defines, fCode);

	glLinkProgram(m_handle);

	if (vAttached)
		glDeleteShader(vAttached);

	if (gAttached)
		glDeleteShader(gAttached);

	if (fAttached)
		glDeleteShader(fAttached);

	int IsLinked = 0;
	glGetProgramiv(m_handle, GL_LINK_STATUS, (int*) &IsLinked);

	if (IsLinked == GL_FALSE)
	{
		int maxLength = 0;
		glGetProgramiv(m_handle, GL_INFO_LOG_LENGTH, &maxLength);

		char* log = new char[maxLength];
		glGetProgramInfoLog(m_handle, maxLength, &maxLength, log);

		__Crash("Failed to compile shader in material %s: \n%s", m_materialName.c_str(), log);
		delete [] log;

		return false;
	}

	/*--------------------------------------------------------------------------*/

	int vs_ubo_bi = AttachUBO("ubo_vs", VERTEX_UBO, EShader::VertexShader);
	int gs_ubo_bi = AttachUBO("ubo_gs", GEOMETRY_UBO, EShader::GeometryShader);
	int fs_ubo_bi = AttachUBO("ubo_fs", FRAGMENT_UBO, EShader::FragmentShader);

	int activeUniforms = 0;
	glGetProgramiv(m_handle, GL_ACTIVE_UNIFORMS, &activeUniforms);

	int uniformNameSize = 0;
	char uniformName[64];

	for (int i = 0; i < activeUniforms; ++i)
	{
		glGetActiveUniformName(m_handle, i, 64, &uniformNameSize, uniformName);

		int blockIndex = 0;
		glGetActiveUniformsiv(m_handle, 1, (uint*) &i, GL_UNIFORM_BLOCK_INDEX, &blockIndex);

		int offset = 0;
		glGetActiveUniformsiv(m_handle, 1, (uint*) &i, GL_UNIFORM_OFFSET, &offset);

		int size = 0;
		glGetActiveUniformsiv(m_handle, 1, (uint*) &i, GL_UNIFORM_SIZE, &size);

		int type = 0;
		glGetActiveUniformsiv(m_handle, 1, (uint*) &i, GL_UNIFORM_TYPE, &type);

		if (type == GL_SAMPLER_2D)
		{
			continue;
		}

		RConst* constant = g_Renderer->GetRendererStorageMgr()->GetConst(std::string(uniformName));
		__Check(constant);

		if (vAttached && blockIndex == vs_ubo_bi)
		{
			m_consts.push_back(RConstantPtn(constant, size, offset, EShader::VertexShader));
			m_sizeofVconsts = std::max(m_sizeofVconsts, offset + constant->GetTotalSize());
			__Check(constant->GetCount() == size);
		}
		else if (gAttached && blockIndex == gs_ubo_bi)
		{
			m_consts.push_back(RConstantPtn(constant, size, offset, EShader::GeometryShader));
			m_sizeofGconsts = std::max(m_sizeofGconsts, offset + constant->GetTotalSize());
			__Check(constant->GetCount() == size);
		}
		else if (fAttached && blockIndex == fs_ubo_bi)
		{
			m_consts.push_back(RConstantPtn(constant, size, offset, EShader::FragmentShader));
			m_sizeofFconsts = std::max(m_sizeofFconsts, offset + constant->GetTotalSize());
			__Check(constant->GetCount() == size);
		}
		else
		{
			__Crash("WAT!?");
		}
	}

	for (auto iter = m_samplers.begin(); iter != m_samplers.end(); ++iter)
	{
		iter->m_location = glGetUniformLocation(m_handle, iter->m_name.c_str());
		__CondError(iter->m_location == -1, "Sampler %s in material %s not found!", iter->m_name.c_str(), m_materialName.c_str());
	}

// 	glBindFragDataLocation(m_programHandle, 0, "FragColor0");
// 	glBindFragDataLocation(m_programHandle, 1, "FragColor1");
// 	glBindFragDataLocation(m_programHandle, 2, "FragColor2");
// 	glBindFragDataLocation(m_programHandle, 3, "FragColor3");

	activeUniforms = 0;

	return true;
}

/*--------------------------------------------------------------------------*/

uint RMaterial::AttachShader( const std::string& macro, GLenum gltype, const std::string& defines, const std::string& code )
{
	if (code.length() > 4)
	{
		std::string src = code;
		std::string samplers;

		PreprocessCode(src);
		src = "#version 330\n#extension GL_ARB_explicit_uniform_location : require\n" + defines + "#define " + macro + " 1\n" + samplers + "\n" + src;

		uint handle = CompileShader(src, gltype);
		__Check(handle);

		glAttachShader(m_handle, handle);
		return handle;
	}

	return 0;
}

/*--------------------------------------------------------------------------*/

int RMaterial::AttachUBO( const std::string& name, uint binding, EShader::Type stage )
{
	int blockIndex = glGetUniformBlockIndex(m_handle, name.c_str());	

	if (blockIndex >= 0)
	{
		glUniformBlockBinding(m_handle, blockIndex, binding);
	}
	
	return blockIndex;
}

/*--------------------------------------------------------------------------*/

uint RMaterial::GetNumSamplers() const
{
	return m_samplers.size();
}

/*--------------------------------------------------------------------------*/

uint RMaterial::GetProgramHandle() const
{
	return m_handle;
}

/*--------------------------------------------------------------------------*/

uint RMaterial::GetLocation( uint i ) const
{
	return m_samplers[i].m_location;
}

/*--------------------------------------------------------------------------*/

const RTexture* RMaterial::GetTexture( uint i ) const
{
	__Check(m_samplers[i].m_texture);
	return m_samplers[i].m_texture;
}

/*--------------------------------------------------------------------------*/

const RSamplerState* RMaterial::GetSamplerState( uint i ) const
{
	return m_samplers[i].m_sampler;
}

/*--------------------------------------------------------------------------*/

void RMaterial::UpdateConstants(uint vOffset, uint gOffset, uint fOffset, RBuffer* buffer) const
{
	uchar* vMem = buffer->Lock(m_sizeofVconsts, vOffset);
	uchar* gMem = buffer->Lock(m_sizeofGconsts, gOffset);
	uchar* fMem = buffer->Lock(m_sizeofFconsts, fOffset);

	for (auto iter = m_consts.begin(); iter != m_consts.end(); ++iter)
	{
		iter->GetRConst()->Update();

		switch(iter->GetStage())
		{
		case EShader::VertexShader:
			{
				memcpy(vMem + iter->GetOffset(), iter->GetRConst()->GetPointer(), iter->GetRConst()->GetTotalSize());
				break;
			}
		case EShader::GeometryShader:
			{
				memcpy(gMem + iter->GetOffset(), iter->GetRConst()->GetPointer(), iter->GetRConst()->GetTotalSize());
				break;
			}
		case EShader::FragmentShader:
			{
				memcpy(fMem + iter->GetOffset(), iter->GetRConst()->GetPointer(), iter->GetRConst()->GetTotalSize());
				break;
			}
		}
	}

	buffer->Flush();

	if (vMem)
	{
		glBindBufferRange(GL_UNIFORM_BUFFER, VERTEX_UBO, buffer->GetHandle(), vOffset, m_sizeofVconsts);
	}

	if (gMem)
	{
		glBindBufferRange(GL_UNIFORM_BUFFER, GEOMETRY_UBO, buffer->GetHandle(), gOffset, m_sizeofGconsts);
	}

	if (fMem)
	{
		glBindBufferRange(GL_UNIFORM_BUFFER, FRAGMENT_UBO, buffer->GetHandle(), fOffset, m_sizeofFconsts);
	}
}

/*--------------------------------------------------------------------------*/

void RMaterial::PreprocessCode( std::string& out )
{
	const char* include = NULL;
	std::string name = "";

	while (include = strstr(out.c_str(), "#include"))
	{
		const char* cpy = include;
		char app = 0;

		while (include[0] != '\n')
		{
			if (include[0] == '"')
			{
				app += 1;
			}
			else if (app == 1)
			{
				name += include[0];
			}

			include++;
		}

		uint beforeStart = 0;
		uint lengthBefore = cpy - out.c_str();

		uint afterStart = include - out.c_str();
		uint lengthAfter = out.size() - afterStart;

		std::string file = ::GetFileContent(DATA_PATH "engine/material/include/" + name);

		std::string before = out.substr(beforeStart, lengthBefore);
		std::string content = "\n// BEGIN OF: " + name + "\n" + file + "\n// END OF " + name + "\n";
		std::string after = out.substr(afterStart, lengthAfter);

		out = before + content + after;
		name = "";
	}
}

/*--------------------------------------------------------------------------*/

uint RMaterial::CompileShader( const std::string& code, GLenum type )
{
	uint handle = glCreateShader(type);

	const char* src = code.c_str();
	int length = code.length();

	glShaderSource(handle, 1, &src, &length);
	glCompileShader(handle);

	int success = GL_FALSE;
	glGetShaderiv(handle, GL_COMPILE_STATUS, &success);

	if (success == GL_FALSE)
	{
		int infoLength = 0;
		glGetShaderiv(handle, GL_INFO_LOG_LENGTH, &infoLength);

		char* log = new char[infoLength];
		glGetShaderInfoLog(handle, infoLength, &success, log);

		__Crash("Failed to compile shader: \n%s", log);
		delete [] log;

		return 0;
	}

	return handle;
}

/*--------------------------------------------------------------------------*/

const std::string& RMaterial::GetMaterialName() const
{
	return m_materialName;
}

/*--------------------------------------------------------------------------*/

const RMaterial* RMaterial::GetNormalVersion() const
{
	return m_normalVersion;
}

/*--------------------------------------------------------------------------*/

const RMaterial* RMaterial::GetInstancedVersion() const
{
	return m_instancedVersion;
}

/*--------------------------------------------------------------------------*/

RSampler::RSampler(const std::string& name, const RSamplerState* sampler, const RTexture* texture, int location ) :
	m_name(name),
	m_sampler(sampler),
	m_texture(texture),
	m_location(-1)
{

}

/*--------------------------------------------------------------------------*/

const std::string& RSampler::GetSamplerName() const
{
	return m_name;
}

/*--------------------------------------------------------------------------*/

const RSamplerState* RSampler::GetSamplerState() const
{
	return m_sampler;
}

/*--------------------------------------------------------------------------*/

const RTexture* RSampler::GetTexture() const
{
	return m_texture;
}

/*--------------------------------------------------------------------------*/

int RSampler::GetLocation() const
{
	return m_location;
}

/*--------------------------------------------------------------------------*/

RMacro::RMacro( const std::string& name, const std::string& value ) :
	m_name(name),
	m_value(value)
{

}

/*--------------------------------------------------------------------------*/

const std::string& RMacro::GetName() const
{
	return m_name;
}

/*--------------------------------------------------------------------------*/

const std::string& RMacro::GetValue() const
{
	return m_value;
}

/*--------------------------------------------------------------------------*/

RConstantPtn::RConstantPtn( const RConst* constant, int size, int offset, EShader::Type stage ) :
	m_offset(offset),
	m_size(size),
	m_stage(stage),
	m_constant(constant)
{

}

/*--------------------------------------------------------------------------*/

int RConstantPtn::GetOffset() const
{
	return m_offset;
}

/*--------------------------------------------------------------------------*/

const RConst* RConstantPtn::GetRConst() const
{
	return m_constant;
}

/*--------------------------------------------------------------------------*/

int RConstantPtn::GetSize() const
{
	return m_size;
}

/*--------------------------------------------------------------------------*/

EShader::Type RConstantPtn::GetStage() const
{
	return m_stage;
}

/*--------------------------------------------------------------------------*/
